# backend/main.py
from fastapi import FastAPI, Request, Response, UploadFile, File
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from typing import List, Dict, Any, Optional
import json, os, time, asyncio, uuid
from datetime import datetime
import threading
import shutil
from medical_analysis import medical_analyzer
import requests
app = FastAPI()
LOCK = threading.Lock()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:5173", "http://127.0.0.1:5173", "*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
os.makedirs(DATA_DIR, exist_ok=True)
SESSIONS_PATH = os.path.join(DATA_DIR, "sessions.json")
MESSAGES_DIR = os.path.join(DATA_DIR, "messages")
os.makedirs(MESSAGES_DIR, exist_ok=True)

def now_iso():
    return datetime.utcnow().isoformat() + "Z"

def load_sessions() -> List[Dict[str, Any]]:
    if not os.path.exists(SESSIONS_PATH):
        with open(SESSIONS_PATH, "w", encoding="utf-8") as f:
            json.dump({"sessions": []}, f, ensure_ascii=False, indent=4)
        return []
    with open(SESSIONS_PATH, "r", encoding="utf-8") as f:
        return json.load(f).get("sessions", [])

def save_sessions(sessions: List[Dict[str, Any]]):
    with open(SESSIONS_PATH, "w", encoding="utf-8") as f:
        json.dump({"sessions": sessions}, f, ensure_ascii=False, indent=4)

def session_messages_path(sid: str):
    return os.path.join(MESSAGES_DIR, f"{sid}.json")

def load_messages(sid: str) -> List[Dict[str, Any]]:
    p = session_messages_path(sid)
    if not os.path.exists(p):
        return []
    with open(p, "r", encoding="utf-8") as f:
        return json.load(f).get("messages", [])

def save_messages(sid: str, messages: List[Dict[str, Any]]):
    with open(session_messages_path(sid), "w", encoding="utf-8") as f:
        json.dump({"messages": messages}, f, ensure_ascii=False, indent=4)

@app.get("/api/sessions")
def list_sessions():
    with LOCK:
        sessions = load_sessions()
    # 统一排序（最近更新在前）
    sessions.sort(key=lambda x: x.get("updated_at", ""), reverse=True)
    return {"sessions": sessions}
@app.post("/api/sessions")
async def create_session(payload: Dict[str, Any] | None = None):
    payload = payload or {}
    sid = str(uuid.uuid4())
    ts = now_iso()
    title = payload.get("title") or f"会话 {sid[:6]}"
    sess = {"id": sid, "title": title, "created_at": ts, "updated_at": ts}
    with LOCK:
      sessions = load_sessions()
      sessions.append(sess)
      save_sessions(sessions)
      save_messages(sid, [])
    return sess

@app.get("/api/sessions/{sid}")
def get_session(sid: str):
    with LOCK:
        msgs = load_messages(sid)
    return {"messages": msgs}

@app.patch("/api/sessions/{sid}")
async def rename_session(sid: str, payload: Dict[str, Any]):
    title = (payload or {}).get("title")
    if not title:
        return JSONResponse({"error":"title required"}, status_code=400)
    with LOCK:
        sessions = load_sessions()
        for s in sessions:
            if s["id"] == sid:
                s["title"] = title
                s["updated_at"] = now_iso()
                break
        save_sessions(sessions)
    return {"ok": True}

@app.delete("/api/sessions/{sid}")
def delete_session(sid: str):
    with LOCK:
        sessions = load_sessions()
        sessions = [s for s in sessions if s["id"] != sid]
        save_sessions(sessions)
        p = session_messages_path(sid)
        if os.path.exists(p):
            os.remove(p)
    return {"ok": True}

@app.post("/api/sessions/{session_id}/update-message")
async def update_message(session_id: str, payload: Dict[str, Any]):
    message_index = payload.get("message_index")
    content = payload.get("content")
    
    if message_index is None or content is None:
        return JSONResponse({"error": "message_index and content required"}, status_code=400)
    
    with LOCK:
        msgs = load_messages(session_id)
        if 0 <= message_index < len(msgs):
            msgs[message_index]["content"] = content
            save_messages(session_id, msgs)
            
            # 更新会话时间
            sessions = load_sessions()
            for s in sessions:
                if s["id"] == session_id:
                    s["updated_at"] = now_iso()
                    break
            save_sessions(sessions)
            
            return {"ok": True}
        return JSONResponse({"error": "message index out of range"}, status_code=400)

@app.post("/api/chat/stream")
async def chat_stream(request: Request, session_id: str, model: str):
    """
    演示用流式输出接口：
    - 接收 messages（包含 user/assistant 历史），将最后一条 user 作为输入
    - 用固定句子分片输出（后续替换为真实模型流即可）
    - SSE: text/event-stream, 每片 "data: xxx\n\n"，最后 "[DONE]"
    """
    body = await request.json()
    messages: List[Dict[str, Any]] = body.get("messages", [])
    user_input = ""
    for m in reversed(messages):
        if m.get("role") == "user":
            user_input = m.get("content", "")
            break

    # 保存消息（把最新 user/assistant 写回）
    with LOCK:
        old = load_messages(session_id)
        # 只追加用户消息（assistant 生成时逐步写入末条即可；演示简单些：生成完再一次性写）
        # 这里为了示例，直接把 last 两条替换为最终结果

    async def sse_generator():
        full_response = ""
        try:
            # 尝试与外部模型服务通信
            with requests.post(
                "http://0.0.0.0:8010/generate_stream",  # 流式接口
                json={"prompt": user_input, "max_new_tokens": 512, "temperature": 0.2, "top_p": 0.95},
                headers={"Authorization": "Bearer mysecret"},
                stream=True  # 关键，启用流式响应
            ) as r:
                r.raise_for_status()
                for chunk in r.iter_content(chunk_size=None, decode_unicode=True):
                    if chunk.startswith("[INST]"):
                        chunk = chunk[len("[INST]"):]
                    if "[/INST]" in chunk:
                        chunk = chunk.split("[/INST]")[-1]
                    if chunk:
                        full_response += chunk
                        yield f"data: {chunk}\n\n"
                        await asyncio.sleep(0.03)
        except requests.exceptions.RequestException as e:
            # 发生错误时返回错误信息
            error_msg = f"请求模型服务出错: {str(e)}"
            yield f"data: {error_msg}\n\n"
            yield "data: [DONE]\n\n"
            return

        # 如果外部服务未返回内容，使用演示回复
        if not full_response:
            # 构建演示回复内容
            demo_response = [
                f"这是针对「{model}」模型的流式响应演示，用于打通前后端。你说的是：{user_input}",
                "",
                "Markdown 支持：**加粗**、`code`、",
                "$$E=mc^2$$",
                "以及列表：",
                "- 项目A",
                "- 项目B",
                "",
                "(后续可替换为真实大模型推理流)。"
            ]
            full_response = "\n".join(demo_response)
            for chunk in split_chunks(full_response, 20):
                yield f"data: {chunk}\n\n"
                await asyncio.sleep(0.03)
        
        yield "data: [DONE]\n\n"

        # 结束后持久化
        with LOCK:
            msgs = load_messages(session_id)
            # 为了让文件里也出现最终 assistant，简单策略：将传入 messages 写为当前状态
            save_messages(session_id, messages)

            # 同时把完整 assistant 回复追加
            msgs = messages[:-1] if messages and messages[-1].get("role") == "assistant" else messages
            msgs.append({"role": "assistant", "content": full_response})
            print(full_response)
            save_messages(session_id, msgs)

            # 更新会话时间
            sessions = load_sessions()
            for s in sessions:
                if s["id"] == session_id:
                    s["updated_at"] = now_iso()
                    break
            save_sessions(sessions)

    headers = {
        "Content-Type": "text/event-stream; charset=utf-8",
        "Cache-Control": "no-cache",
        "Connection": "keep-alive",
        "X-Accel-Buffering": "no",  # Nginx 反代时建议
    }
    return StreamingResponse(sse_generator(), headers=headers)

# @app.post("/api/chat/stream")
# async def chat_stream(request: Request, session_id: str, model: str):
#     """
#     演示用流式输出接口：
#     - 接收 messages（包含 user/assistant 历史），将最后一条 user 作为输入
#     - 用固定句子分片输出（后续替换为真实模型流即可）
#     - SSE: text/event-stream, 每片 "data: xxx\n\n"，最后 "[DONE]"
#     """
#     body = await request.json()
#     messages: List[Dict[str, Any]] = body.get("messages", [])
#     user_input = ""
#     for m in reversed(messages):
#         if m.get("role") == "user":
#             user_input = m.get("content", "")
#             break

#     # 固定回复：分片输出
#     demo_reply = f"这是针对「{model}」模型的流式响应演示，用于打通前后端。你说的是：{user_input}\n\n" \
#                  "Markdown 支持：**加粗**、`code`、\n" \
#                  "$$E=mc^2$$\n" \
#                  "以及列表：\n- 项目A\n- 项目B\n\n(后续可替换为真实大模型推理流)。"

#     async def sse_generator():
#         # 保存消息（把最新 user/assistant 写回）
#         with LOCK:
#             old = load_messages(session_id)  # pyright: ignore[reportUnusedVariable]
#             # 只追加用户消息（assistant 生成时逐步写入末条即可；演示简单些：生成完再一次性写）
#             # 这里为了示例，直接把 last 两条替换为最终结果
#         # 流式发数据
#         for chunk in split_chunks(demo_reply, 20):
#             yield f"data: {chunk}\n\n"
#             await asyncio.sleep(0.03)
#         yield "data: [DONE]\n\n"

#         # 结束后持久化
#         with LOCK:
#             msgs = load_messages(session_id)
#             # 为了让文件里也出现最终 assistant，简单策略：将传入 messages 写为当前状态
#             save_messages(session_id, messages)

#             # 同时把完整 assistant 回复追加
#             msgs = messages[:-1] if messages and messages[-1].get("role") == "assistant" else messages
#             msgs.append({"role": "assistant", "content": demo_reply})
#             save_messages(session_id, msgs)

#             # 更新会话时间
#             sessions = load_sessions()
#             for s in sessions:
#                 if s["id"] == session_id:
#                     s["updated_at"] = now_iso()
#                     break
#             save_sessions(sessions)

#     headers = {
#         "Content-Type": "text/event-stream; charset=utf-8",
#         "Cache-Control": "no-cache",
#         "Connection": "keep-alive",
#         "X-Accel-Buffering": "no",  # Nginx 反代时建议
#     }
#     return StreamingResponse(sse_generator(), headers=headers)

def split_chunks(s: str, n: int):
    return [s[i:i+n] for i in range(0, len(s), n)]

# ===================== 医学影像分析API ，最后更新时间2025-09-08 11:00=====================
MEDICAL_UPLOADS_DIR = os.path.join(DATA_DIR, "medical_uploads")
os.makedirs(MEDICAL_UPLOADS_DIR, exist_ok=True)

@app.post("/api/medical/upload-image")
async def upload_medical_image(file: UploadFile = File(...)):
    """上传医学影像文件"""
    try:
        # 检查文件类型
        if not file.content_type.startswith('image/'):
            return JSONResponse(
                {"error": "只支持图像文件"}, 
                status_code=400
            )
        
        # 生成唯一文件名
        timestamp = int(time.time())
        filename = f"{timestamp}_{file.filename}"
        file_path = os.path.join(MEDICAL_UPLOADS_DIR, filename)
        
        # 保存文件
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        return {
            "success": True,
            "filename": filename,
            "file_path": file_path,
            "message": "文件上传成功"
        }
        
    except Exception as e:
        return JSONResponse(
            {"error": f"文件上传失败: {str(e)}"}, 
            status_code=500
        )

@app.post("/api/medical/analyze")
async def analyze_medical_image(image: UploadFile = File(...), custom_prompt: str = None):
    """分析医学影像"""
    try:
        # 保存上传的图像文件
        if not image.filename:
            return JSONResponse(
                {"error": "未选择图像文件"}, 
                status_code=400
            )
        
        # 创建上传目录
        upload_dir = os.path.join(DATA_DIR, "medical_uploads")
        os.makedirs(upload_dir, exist_ok=True)
        
        # 生成唯一文件名
        file_extension = os.path.splitext(image.filename)[1]
        filename = f"analysis_{int(time.time())}{file_extension}"
        file_path = os.path.join(upload_dir, filename)
        
        # 保存文件
        with open(file_path, "wb") as buffer:
            content = await image.read()
            buffer.write(content)
        
        # 执行医学影像分析
        result = medical_analyzer.analyze_medical_image(file_path, custom_prompt)
        
        # 清理临时文件
        try:
            os.remove(file_path)
        except:
            pass
        
        return result
        
    except Exception as e:
        return JSONResponse(
            {"error": f"分析失败: {str(e)}"}, 
            status_code=500
        )

@app.post("/api/medical/analyze-complete")
async def analyze_medical_image_complete(
    image: UploadFile = File(...), 
    input_text: str = None,
    temperature: float = 0.8,
    top_p: float = 0.4
):
    """完整的医学影像分析（基于demo4.py）"""
    try:
        # 保存上传的图像文件
        if not image.filename:
            return JSONResponse(
                {"error": "未选择图像文件"}, 
                status_code=400
            )
        
        # 创建上传目录
        upload_dir = os.path.join(DATA_DIR, "medical_uploads")
        os.makedirs(upload_dir, exist_ok=True)
        
        # 生成唯一文件名
        file_extension = os.path.splitext(image.filename)[1]
        filename = f"complete_analysis_{int(time.time())}{file_extension}"
        file_path = os.path.join(upload_dir, filename)
        
        # 保存文件
        with open(file_path, "wb") as buffer:
            content = await image.read()
            buffer.write(content)
        
        # 执行完整的医学影像分析
        result = medical_analyzer.analyze_medical_image_complete(
            file_path, 
            input_text, 
            temperature, 
            top_p
        )
        
        # 清理临时文件
        try:
            os.remove(file_path)
        except:
            pass
        
        return result
        
    except Exception as e:
        return JSONResponse(
            {"error": f"分析失败: {str(e)}"}, 
            status_code=500
        )

@app.post("/api/medical/continue-conversation")
async def continue_medical_conversation(request: Request):
    """继续医学影像对话（后续问答使用本地模型或API）"""
    try:
        body = await request.json()
        image_path = body.get("image_path")
        input_text = body.get("input_text")
        history = body.get("history", [])
        temperature = body.get("temperature", 0.8)
        top_p = body.get("top_p", 0.4)
        
        if not image_path or not input_text:
            return JSONResponse(
                {"error": "图像路径和输入文本不能为空"}, 
                status_code=400
            )
        
        # 调用后续对话函数
        result = medical_analyzer.continue_medical_conversation(
            image_path, input_text, history, temperature, top_p
        )
        
        return result
        
    except Exception as e:
        return JSONResponse(
            {"error": f"对话失败: {str(e)}"}, 
            status_code=500
        )

@app.post("/api/medical/update-patient")
async def update_patient_info(request: Request):
    """更新患者信息"""
    try:
        body = await request.json()
        name = body.get("name")
        age = body.get("age")
        gender = body.get("gender")
        history = body.get("history")
        
        result = medical_analyzer.update_patient_info(name, age, gender, history)
        
        return result
        
    except Exception as e:
        return JSONResponse(
            {"error": f"更新患者信息失败: {str(e)}"}, 
            status_code=500
        )

@app.get("/api/medical/patient-info")
async def get_patient_info():
    """获取当前患者信息"""
    return {
        "success": True,
        "patient_info": medical_analyzer.patient_info.copy()
    }

@app.post("/api/medical/generate-report")
async def generate_medical_report(request: Request):
    """生成医疗报告"""
    try:
        body = await request.json()
        analysis_history = body.get("analysis_history", [])
        chat_history = body.get("chat_history", [])
        patient_info = body.get("patient_info", {})
        
        # 如果有聊天历史，使用聊天历史生成报告
        if chat_history:
            report = medical_analyzer.generate_medical_report_from_chat(chat_history, patient_info)
        else:
            # 否则使用分析历史
            report = medical_analyzer.generate_medical_report(analysis_history)
        
        return {
            "success": True,
            "report": report,
            "timestamp": datetime.now().isoformat()
        }
        
    except Exception as e:
        return JSONResponse(
            {"error": f"报告生成失败: {str(e)}"}, 
            status_code=500
        )

# ===================== 百度翻译API ，最后更新时间2025-10-18 =====================
import random
from hashlib import md5

# Set your own appid/appkey.
appid = '20251013002474387'
appkey = 'sbdaZJn0eKEdwloyrhHe'

# For list of language codes, please refer to `https://api.fanyi.baidu.com/doc/21`
from_lang = 'en'
to_lang = 'zh'

endpoint = 'http://api.fanyi.baidu.com'
path = '/api/trans/vip/translate'
url = endpoint + path

# Generate salt and sign
def make_md5(s, encoding='utf-8'):
    return md5(s.encode(encoding)).hexdigest()

@app.post("/api/translate")
async def translate_text(request: Request):
    """调用百度翻译API进行文本翻译"""
    try:
        body = await request.json()
        query = body.get("text", "")
        
        if not query:
            return JSONResponse(
                {"error": "文本内容不能为空"}, 
                status_code=400
            )
        # 检查是否包含代码块
        query_parts = query.split("```")
        print(f"分割后的部分数量: {len(query_parts)}")
        has_code_block = len(query_parts) > 1
        
        # 如果有代码块，只翻译非代码块部分
        if has_code_block:
            translated_parts = []
            for i, part in enumerate(query_parts):
                # 代码块内容（奇数索引位置）不翻译
                if i % 2 == 1:
                    translated_parts.append(part)
                # 非代码块内容（偶数索引位置）需要翻译
                else:
                    if part.strip():
                        print(f"翻译第{i}部分: {part[:30]}...")
                        # 为每个需要翻译的部分生成salt和sign
                        salt = random.randint(32768, 65536)
                        sign = make_md5(appid + part + str(salt) + appkey)
                        
                        # 构建请求
                        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
                        payload = {'appid': appid, 'q': part, 'from': from_lang, 'to': to_lang, 'salt': salt, 'sign': sign}
                        
                        # 发送请求
                        r = requests.post(url, params=payload, headers=headers)
                        result = r.json()
                        print(f"翻译结果: {result}")
                        
                        # 检查是否有错误
                        if 'error_code' in result:
                            print(f"翻译错误: {result['error_code']}")
                            # 翻译失败时保留原文
                            translated_parts.append(part)
                        elif 'trans_result' in result and len(result['trans_result']) > 0:
                            # 拼接所有翻译结果
                            translated_text_part = '\n\n'.join([item['dst'] for item in result['trans_result']])
                            translated_parts.append(translated_text_part)
                            # translated_parts.append(result['trans_result'][0]['dst'])
                        else:
                            translated_parts.append(part)
                    else:
                        translated_parts.append(part)
            
            # 组合翻译后的内容，保持原始格式
            translated_text = '\n```\n'.join(translated_parts)
            print(f"最终翻译结果: {translated_text[:100]}...")
            return {
                "success": True,
                "result": translated_text,
                "timestamp": datetime.now().isoformat()
            }
        # Generate salt and sign
        salt = random.randint(32768, 65536)
        sign = make_md5(appid + query + str(salt) + appkey)
        
        # Build request
        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
        payload = {'appid': appid, 'q': query, 'from': from_lang, 'to': to_lang, 'salt': salt, 'sign': sign}
        
        # Send request
        r = requests.post(url, params=payload, headers=headers)
        result = r.json()
        print(result)
        
        # 检查是否有错误
        if 'error_code' in result:
            return JSONResponse(
                {"error": f"翻译失败: {result.get('error_msg', '未知错误')}"}, 
                status_code=500
            )
        
        # 直接返回dst字段的翻译结果
        if 'trans_result' in result and len(result['trans_result']) > 0:
            return {
                "success": True,
                "result": result['trans_result'][0]['dst'],  # 直接返回dst字段
                "timestamp": datetime.now().isoformat()
            }
        else:
            return JSONResponse(
                {"error": "翻译结果格式异常"}, 
                status_code=500
            )
        
    except Exception as e:
        return JSONResponse(
            {"error": f"翻译请求失败: {str(e)}"}, 
            status_code=500
        )
